#! /usr/bin/env python
# -*- coding: utf-8 -*-

# *************************************************************
#  Copyright (c) Huoty - All rights reserved
#
#      Author: Huoty <sudohuoty@gmail.com>
#  CreateTime: 2016-12-26 12:55:28
# *************************************************************

from __future__ import print_function

import sys
import signal
from resource import getrlimit, setrlimit, RLIMIT_DATA, RLIMIT_STACK, RLIMIT_AS, RLIMIT_RSS

try:
    from resource import RLIMIT_VMEM
except ImportError:
    pass

def handle_sigsegv(signum, frame):
    print(signal, frame)
    sys.exit(1)

def limitmem(resource, limit):
    resource = {
        "heap": RLIMIT_DATA,
        "stack": RLIMIT_STACK,
        "vm": RLIMIT_VMEM if sys.version_info[0] > 2 else None,
        "as": RLIMIT_AS,
        "rss": RLIMIT_RSS,
    }.get(resource, None)

    if not resource:
        raise Exception("No limit target `{}`".format(resource))

    limit = limit * 1024 * 1024 * 1024

    print("set before resource:", getrlimit(resource))
    setrlimit(resource, (limit, limit))
    print("set after resource:", getrlimit(resource))

def consume_mem():
    dct = {}
    from random import random
    while 1:
        dct.update({random(): random()})
    pass


# Script starts from here

if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('-H', action='store_true', help="limit heap")
    parser.add_argument('-s', action='store_true', help="limit stack")
    parser.add_argument('-v', action='store_true', help="limit vm")
    parser.add_argument('-a', action='store_true', help="limit as")
    parser.add_argument('-r', action='store_true', help="limit rss")

    options = parser.parse_args()

    if options.H:
        resource = "heap"
    elif options.s:
        resource = "stack"
    elif options.v and sys.version_info[0] > 2:
        resource = "vm"
    elif options.a:
        resource = "as"
    elif options.r:
        resource = "rss"
    else:
        parser.print_help()
        sys.exit(0)

    limitmem(resource, 0.00000001)

    # 注册系统信号
    signal.signal(signal.SIGSEGV, handle_sigsegv)

    consume_mem()
